import os
import logging
import logging.config
import FileUtility
from datamodel import HmmerHit, DomainMatch, HmmerHitValidation, HomologFinderConfig, NoSequencesToExtractError
from Bio import AlignIO, SeqIO
from Bio.Blast import NCBIStandalone
import re
import glob
import random
import urllib
import FileUtility

logging.config.fileConfig("logging.conf")
logger = logging.getLogger("MAIN_LOG")


def makeStockholmAlignment(alignment, outputFile):
    if(not FileUtility.safeFileExists(outputFile)):
        logger.debug("Making a stockholm alignment from a fasta alignment - outputting to %s" % outputFile)
        inputHandle = None
        outputHandle = None
        try:
            inputHandle = open(alignment, "rU", 1)
            outputHandle = open(outputFile, "w")
            alignments = AlignIO.parse(inputHandle, "fasta")
            AlignIO.write(alignments, outputHandle, "stockholm")
        finally:
            FileUtility.safeFileClose(inputHandle)
            FileUtility.safeFileClose(outputHandle)


def makeHmm(alignment, outputFile):
    if(not FileUtility.safeFileExists(outputFile)):
        logger.debug("Making an hmm from %s to %s" % (alignment, outputFile))
        str = HomologFinderConfig.HMMER_BUILD + " %s %s" % (outputFile, alignment)
        os.system(str)

def runHmmer(sequences, hmmFile, outputFile):
    if(not FileUtility.safeFileExists(outputFile)):
        logger.debug("Searching for homologs using %s in %s - outputting to %s" % (hmmFile, sequences, outputFile))
        str = HomologFinderConfig.HMMER_SEARCH + " -o %s %s %s" % (outputFile, hmmFile, sequences)
        os.system(str)


def parseHmmerOutput(input):
    file = None
    map = None
    try:
        file = open(input, 'r', 1) # 1 means that it's buffered
        hit=False
        map = {HmmerHitValidation.HOMOLOGOUS : list(), HmmerHitValidation.MULTIDOMAIN : list(), HmmerHitValidation.REPETITIOUS : list()}
        for line in file:
            if(hit):
                #There may not be an inclusion threshold, it may just end with a blank line
                if (re.search("inclusion threshold", line) or line=='\n'):
                    hit=False
                    break
                else:
                    hmmerHit = HmmerHit(line)
                    validation = hmmerHit.validate()
                    map[validation].append(hmmerHit)
            elif(re.search("--- full sequence ---", line)):
                next(file)
                next(file)
                hit = True
    finally:
        FileUtility.safeFileClose(file)
    return map


def createBlastDb(fastaFile, dbName, protein=True):
    '''Create a database for blasting analysis

    This method checks to see if a file is amino acids or nucleotides
    '''
    if (not glob.glob(dbName + '.*')):
        logger.info("Creating database for %s at %s" % (fastaFile, dbName))
        str = HomologFinderConfig.BLAST_FORMATDB + " -p %s -o T -i %s -n %s" % ('T' if protein else 'F', fastaFile, dbName)
        os.system(str)
    else:
        #logger.debug('Database %s already exists' % dbName)
        pass



def getHmmFromPfam(hmmId, outputFile):
    """ Uses hmmsearch to return an hmm from the pfam-A database"""
    if(not FileUtility.safeFileExists(outputFile)):
        try:
            urllib.urlretrieve(HomologFinderConfig.PFAM_HOME_URL + hmmId, outputFile)
        except:
            logger.error("Error retrieving %s from the PFAM database at %s" % (hmmId, HomologFinderConfig.PFAM_HOME_URL))
    else:
        logger.debug("%s has already been retrieved (stored at %s)" %(hmmId, outputFile))

def makeConsensusFileFromHmm(hmmFile, consensusFile):
    """Creates a consensus file from an HMM file using hmmemit"""
    if(not FileUtility.safeFileExists(consensusFile)):
        str = HomologFinderConfig.HMMER_EMIT + " -c -o %s %s" % (consensusFile, hmmFile)
        os.system(str)

def writeFastaFile(fastas, outputFile):
    """writes the fastas to a file"""
    if(not FileUtility.safeFileExists(outputFile) and len(fastas)):
        try:
            handle = open(outputFile, "w", 1)
            SeqIO.write(fastas, handle, "fasta")
        finally:
            FileUtility.safeFileClose(handle)

def getSequencesFromDb(sequenceList, db, output=None):
    '''returns a list of sequences from a database using fastacmd
            if output is not set, the method returns a list of fastas
            otherwise it returns a filename linking to the fastas
        Throws an exception if no sequences are passed in
    '''
    if (not len(sequenceList)):
        #raise NoSequencesToExtractError("There were no sequences to extract from %s" % db)
        logger.error("There were no sequences to extract from %s" % db)
        return

    tmpFile = HomologFinderConfig.SCRATCH + '%016x.tmp' % random.getrandbits(128)
    outputFile = HomologFinderConfig.SCRATCH + '%016x.tmp' % random.getrandbits(128) if output is None else output
    #write out the list to a file
    f = open(tmpFile, 'w', 1)
    for seq in sequenceList:
        f.write(seq + "\n")
    f.close()
    str = HomologFinderConfig.BLAST_FASTACMD + " -d %s -i %s -o %s" % (db, tmpFile, outputFile)
    os.system(str)
    FileUtility.safeFileRemove([tmpFile])

    if output is None:
        fastaList =list()
        index = SeqIO.index(outputFile, 'fasta')
        for key in index.keys():
            fastaList.append(index[key])
        FileUtility.safeFileRemove([outputFile])
        return fastaList
    else:
        return output


def getHomologListFromBLASTOutput(blastOut):
    '''returns a list of sequences from parsing a blast output file'''

    matches = set() # keeping this a set will remove identical records

    if (os.path.getsize(blastOut)):

        try:
            handle = open(blastOut, 'r')
            blastParser = NCBIStandalone.BlastParser()
            blastRecord = blastParser.parse(handle)
            for alignment in blastRecord.alignments:
                match = re.match('>(\S+)', alignment.title)
                matches.add(match.group(1))
            handle.close()
        except:
            logger.error("Something went wrong while parsing the file %s, it may just be empty" % blastOut)
    return matches

def isSearchPickled(hs):
    ''' Checks to see if a homolog search has been preserved from a previous run'''
    return FileUtility.safeFileExists(HomologFinderConfig.PICKLE_STORE + hs.storage)


def pickleSearch(hs):
    '''Stores a homolog search to the pickle store
        This file can be parsed using a grep on >> to pull out the count record'''


    handle = None
    try:
        handle = open(HomologFinderConfig.PICKLE_STORE + hs.storage, 'w', 1)
                                                              
        handle.write('# Report for %s (%s) vs %s \n\n' % (hs.hik.geneName, hs.hik.cog, hs.db.name))
        handle.write('#Gene\tCOG\tdatabase\tgood.hits\tbad.hits\tgood.species.count\tbad.species.count\n')
        tmp = set()
        for hit in hs.goodHits:
            tmp.add(hit.species)
        goodSpeciesCount = len(tmp)
        tmp = set()
        for hit in hs.badHits:
            tmp.add(hit.species)
        badSpeciesCount = len(tmp)
        handle.write('>>' + '\t'.join([hs.hik.geneName, hs.hik.cog, hs.db.name, str(len(hs.goodHits)), str(len(hs.badHits)), str(goodSpeciesCount), str(badSpeciesCount)]) + '\n\n')
        handle.write('--- Homologs which passed reciprocal best hit ---\n')
        handle.write('#fragment\tbest.string.hit\tspecies\n')
        for hit in hs.goodHits:
            handle.write('%s\t%s\t%s\n' % (hit.fragment, hit.stringHit, hit.species))
        handle.write('\n\n--- Homologs which failed reciprocal best hit ---\n')
        handle.write('#fragment\tbest.string.hit\tspecies\n')
        for hit in hs.badHits:
            handle.write('%s\t%s\t%s\n' % (hit.fragment, hit.stringHit, hit.species))
    finally:
        FileUtility.safeFileClose(handle)


def unpickleSearch(hs):
    '''Unpickles a search from the store'''
    #todo write this method

if __name__ == "__main__":
    #fastas = getSequencesFromDb(['JCVI_READ_564873', 'JCVI_READ_338360', 'JCVI_READ_624694', 'JCVI_READ_151370'], "/data/scratch/GStest", output='/data/scratch/test.txt')
    hits = getHomologListFromBLASTOutput('/data/pipeline/scratch/GStest.ppk1.homologs.vs.STRING.out')
